#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# A tool for transferring data between two peers, like nc(netcat) utility.
#

import sys
import socket
import argparse
import traceback

def args_parser():
    comment = ("A tool for transferring data between two peers, like nc(netcat) utility.")
    epilog = ("Examples:\n"
              "\tServer: %(prog)s -s -t 10\n"
              "\tClient: %(prog)s -c\n"
              "\n"
              "\tServer: %(prog)s -s -r\n"
              "\tClient: %(prog)s -c -w -m hello\n"
              "\tClient: echo hello | %(prog)s -c -w\n"
              "\n"
              "\tServer: %(prog)s -s -w -m hello\n"
              "\tServer: echo hello | %(prog)s -s -w\n"
              "\tClient: %(prog)s -c -r\n"
              "\n"
              "\tServer: %(prog)s -s -rw -m hello\n"
              "\tClient: %(prog)s -c -wr -m hello\n"
              "\n"
              "\tServer: %(prog)s -s -r > output\n"
              "\tClient: %(prog)s -c -w < input\n"
              "\n"
              "\tServer: %(prog)s -s -b 192.168.2.104 -p 8864 -r\n"
              "\tClient: %(prog)s -c 192.168.2.104 -p 8864 -w -m hello\n"
              "\tClient: %(prog)s -c 192.168.2.104 -p 8864 -b 192.168.2.106 -w -m hello\n"
              "\n")
    parser = argparse.ArgumentParser(description = comment, epilog = epilog,
                                     formatter_class = argparse.RawDescriptionHelpFormatter)
    parser.add_argument('-s', '--server', dest = 'server', action = 'store_true',
                        help = "Run in server mode.")
    parser.add_argument('-c', '--client', dest = 'serverhost', action = 'store',
                        nargs = '?', metavar = 'host', default = None,
                        help = "Run in client mode and connect to <host>.")
    parser.add_argument('-b', '--bind', dest = 'localhost', action = 'store',
                        metavar = 'host', default = None,
                        help = "Bind to <host>, an interface address.")
    parser.add_argument('-p', '--port', dest = 'port', action = 'store',
                        metavar = 'port', default = 9981,
                        help = "Set server port to listen on/connect to <port>.")
    parser.add_argument('-r', '--read', dest = 'read', action = 'store_true',
                        help = "Read data from peer.")
    parser.add_argument('-w', '--write', dest = 'write', action = 'store_true',
                        help = "Write data to peer.")
    parser.add_argument('-t', '--timeout', dest = 'timeout', action = 'store',
                        metavar = 'timeout', type = float, default = None,
                        help = "Exit after idle <timeout> seconds when accepting connections, "
                               "establishing connections or receiving data, etc.")
    parser.add_argument('-6', '--IPv6', dest = 'IPv6', action = 'store_true',
                        help = "Set the domain to IPv6 instead of the default option of IPv4.")
    parser.add_argument('-m', '--message', dest = 'message', action = 'store',
                        nargs = '*', metavar = 'message', default = None,
                        help = "Send <message> instead of reading from stdin.")
    parser.add_argument('-v', '--verbose', dest = 'verbose', action='store_true',
                        help = "Print detail log to stderr.")
    parser.add_argument('-S', '--size', dest = 'size', action='store',
                        metavar = '1-1073741824', type = int, default = 1024,
                        help = "Read/Write at most <size> bytes each time.")
    return parser

def usage():
    args_parser().parse_args(["-h"])

def parse_args(args_list = None):
    args = args_parser().parse_args(args_list)
    if args.size < 1 or args.size > 1073741824:
        args.size = 1024
    return args

def format_exception():
    res = ''
    exc_type, exc_value, exc_traceback = sys.exc_info()
    estr = traceback.format_exception(exc_type, exc_value, exc_traceback)
    for s in estr:
        res += s
    return res

def print_verbose(string):
    if args.verbose:
        # write all log to stderr
        sys.stderr.write(string)

def recv_data(sock):
    if args.read:
        print_verbose(f"Receive: {sock}\n")
        total = b''
        data = sock.recv(args.size)
        while data:
            # write all received data to stdout
            sys.stdout.buffer.write(data)
            total += data
            data = sock.recv(args.size)
        print_verbose(f"Received: {total}\n")

def send_data(sock):
    if not args.write:
        sock.shutdown(socket.SHUT_WR)
        print_verbose(f"Shutdown write: {sock}\n")
        return
    if args.message:
        print_verbose(f"Send message: {sock}\n")
        for msg in args.message:
            sock.send(msg.encode())
            print_verbose(f"Sent: {msg}\n")
    else:
        print_verbose(f"Read stdin and send: {sock}\n")
        total = b''
        data = sys.stdin.buffer.read(args.size)
        while data:
            sock.send(data)
            total += data
            data = sys.stdin.buffer.read(args.size)
        print_verbose(f"Sent: {total}\n")
    sock.shutdown(socket.SHUT_WR)
    print_verbose(f"Shutdown write: {sock}\n")

def server_run():
    try:
        family = socket.AF_INET6 if args.IPv6 else socket.AF_INET
        addrinfo = socket.getaddrinfo(args.localhost, args.port,
                                      family, socket.SOCK_STREAM,
                                      0, socket.AI_PASSIVE)
        family, socktype, proto, canonname, localaddr = addrinfo[0]
        sock = socket.socket(family, socktype, proto)
        sock.settimeout(args.timeout)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(localaddr)
        sock.listen()
        print_verbose(f'Listen: {sock}\n')
        conn, addr = sock.accept()
        try:
            print_verbose(f'Connected: {conn}\n')
            recv_data(conn)
            send_data(conn)
        except:
            print_verbose(f"{format_exception()}\n")
            sys.exit(-1)
        finally:
            print_verbose(f'Close: {conn}\n')
            conn.close()
    except:
        print_verbose(f"{format_exception()}\n")
        sys.exit(-1)
    finally:
        print_verbose(f'Close: {sock}\n')
        sock.close()

def client_run():
    try:
        family = socket.AF_INET6 if args.IPv6 else socket.AF_INET
        addrinfo = socket.getaddrinfo(args.serverhost, args.port,
                                      family, socket.SOCK_STREAM,
                                      0, socket.AI_PASSIVE)
        family, socktype, proto, canonname, serveraddr = addrinfo[0]
        addrinfo = socket.getaddrinfo(args.localhost, 0,
                                      family, socket.SOCK_STREAM,
                                      0, socket.AI_PASSIVE)
        family, socktype, proto, canonname, localaddr = addrinfo[0]
        sock = socket.socket(family, socktype, proto)
        sock.settimeout(args.timeout)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(localaddr)
        print_verbose(f'Connect: {sock} -> {serveraddr}\n')
        sock.connect(serveraddr)
        print_verbose(f'Connected: {sock}\n')
        send_data(sock)
        recv_data(sock)
    except:
        print_verbose(f"{format_exception()}\n")
        sys.exit(-1)
    finally:
        print_verbose(f'Close: {sock}\n')
        sock.close()

args = parse_args()
print_verbose(f"{args}\n")
if args.server:
    server_run()
else:
    client_run()
